#include "../Base.h"


#ifndef _WIN64
PUCHAR gMapNtdll = NULL;

BOOLEAN SsdtInit()
{
	if (gMapNtdll) return NULL;
	WCHAR* path = GetSystemRootNtdllPath();
	gMapNtdll =(PUCHAR) MapOfSection(path);
	ExFreePoolWithTag(path, 'yang');
	return TRUE;

}


VOID SsdtDestory()
{
	if(!gMapNtdll)
	UnMapOfSection(gMapNtdll);
}

ULONG SsdtGetFunctionIndex(char* funcName)
{

	ULONG_PTR funcAddress;
	funcAddress = ExportTableFuncByName((char*)gMapNtdll, funcName);
	if (!funcAddress) return -1;

	return *(PULONG)((PUCHAR)funcAddress+1);
}

__declspec(naked) ULONG wpOff()
{

	_asm {
		sti;
		push ebx;
		mov eax, cr0;
		mov ebx, 10000;
		not ebx;
		and eax, ebx;
		mov cr0, eax;
		pop ebx;
		ret;
	}

}

__declspec(naked) VOID wpOn(ULONG cr0)
{
	_asm{
		mov eax, dword ptr ds : [esp + 8] ;
		cli;
	}
}

SsdtServiceTableEntry SsdtGetTable()
{
	return KeServiceDescriptorTable->Ssdt;
}


ULONG_PTR SsdtSetHook(char* funcName,PVOID newFuncAddress)
{
	ULONG funAddress;
	SsdtServiceTableEntry Ssdt;
	PULONG SsdtTableBase;
	
	Ssdt = SsdtGetTable();
	SsdtTableBase = (PULONG)Ssdt.funTable;
	gMapNtdll = (PUCHAR)MapOfSection(GetSystemRootNtdllPath());
	
	ULONG index = SsdtGetFunctionIndex("NtOpenProcess");
	
	if (index == -1) return 0;
	funAddress = SsdtTableBase[index];
	ULONG wpFlag = wpOff();
	SsdtTableBase[index] = (ULONG)newFuncAddress;
	wpOn(wpFlag);
	return funAddress;
}



ULONG SsdtGetFunctionAddress(char* funcName) {

	SsdtServiceTableEntry Ssdt;
	PULONG SsdtTableBase;
	ULONG funcAddress;

	Ssdt = SsdtGetTable();
	SsdtTableBase = (PULONG)Ssdt.funTable;
	gMapNtdll = (PUCHAR)MapOfSection(GetSystemRootNtdllPath());
	ULONG index = SsdtGetFunctionIndex(funcName);
	if (index == -1) return 0;
	funcAddress = SsdtTableBase[index];

	return funcAddress;
}

#else



#endif